
from app.db.models.asset_tasks import get_monitor_urls

from dateutil import parser
from . import BaseHandler
from ...libs.score import Score

from .utils import get_warning_id
from .events import Event


VUL_LEVEL_MAP = {
    "高": 4,
    "中": 3,
    "低": 2,
}


class SslHandler(BaseHandler):
    def handle(self):
        return self._handle_ssl(task_settings=self.task_settings, response=self.response, target=self.target)

    def _handle_ssl(self, task_settings, response, target):
        results = response.get("sslcheck", {}).get("results", [])
        data = response.get("sslcheck", {}).get("data", {})
        plugin_categories = task_settings.get("collect", {}).get("pluginCategories", [])
        warnings = []
        for item in results:
            item["foundAt"] = parser.parse(item.pop("foundAt")) if item.get("foundAt") else None
            category = item.get("category")
            if category not in plugin_categories:
                continue
            warning = {
                "affects": target,
                "title": item.get("display_name"),
                "description": item.get("description"),
                "category": item.get("category"),
                "impact": item.get("impact"),
                "recommendation": item.get("recommendation"),
                "detail": item,
                "id": get_warning_id(
                    risk_title=item.get("name"),
                    task_type="ssl",
                    url=target,
                    job_id=self.task.jobId,
                ),
                "level": VUL_LEVEL_MAP.get(item.get("risk")) or item.get("risk"),
                "levelName": Score.severityMap.get(item.get("risk"), {}).get("zh_CN") or item.get("risk"),
                "foundAt": response.get("found_at"),
                "detailText": item.get("extra_detail", ""),
                "isResolved": False
            }
            warnings.append(warning)

        statistics = data.get("statistics") or {}
        statistics.update({
            "urls": len(get_monitor_urls(task_session_id=self.task.taskSessionId, is_source_data=True)),
            "scan_count": get_monitor_urls(task_session_id=self.task.taskSessionId, is_scan_count=True, is_source_data=True),
        })
        # 防御信息
        data.update({
            "site_info": {"waf_list": response.get("waf_list", [])},
            "target_ip": response.get("target_ip", ""),
            "statistics": statistics
        })

        warnings_addition = {
            "addition": data,
            "warnings": warnings,
        }
        if target_status := Event._target_status(response):
            warnings_addition["targetStatus"] = target_status
        return warnings_addition
